#!/usr/bin/env python3
"""
ML预测功能测试脚本
用于验证机器学习模型预测功能是否正常工作
"""

import sys
import os
import numpy as np

# 添加路径以便导入模块
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(os.path.dirname(__file__))

try:
    from utils.ml_predictor import predict_ml_comm_time, get_ml_predictor
    from ML_analyze import get_ML_Prediction_results
    print("✅ 成功导入ML预测模块")
except ImportError as e:
    print(f"❌ 导入ML预测模块失败: {e}")
    sys.exit(1)

def test_ml_predictor():
    """测试ML预测器基本功能"""
    print("\n🔬 测试ML预测器基本功能...")
    
    try:
        # 测试不同消息大小的预测
        test_sizes = [26504, 16384, 65536, 262144, 1048576]  # 1KB, 16KB, 64KB, 256KB, 1MB
        
        print("\n📊 测试不同消息大小的预测结果:")
        print("消息大小(字节)\t预测延迟(μs)")
        print("-" * 40)
        
        for size in test_sizes:
            try:
                # 测试平均预测（4种通信类型的平均值）
                avg_prediction = predict_ml_comm_time(size, comm_type=None)
                print(f"{size:>12}\t{avg_prediction:>10.3f}")
                
                # 验证预测值是否合理（应该为正数）
                assert avg_prediction > 0, f"预测值应该为正数，但得到: {avg_prediction}"
                
            except Exception as e:
                print(f"❌ 预测消息大小 {size} 时出错: {e}")
                return False
        
        print("✅ ML预测器基本功能测试通过")
        return True
        
    except Exception as e:
        print(f"❌ ML预测器测试失败: {e}")
        return False

def test_individual_comm_types():
    """测试各个通信类型的预测"""
    print("\n🔬 测试各个通信类型的预测...")
    
    try:
        message_size = 65536  # 64KB
        comm_types = [1, 2, 3, 4]
        type_names = ["Intra-Node", "Intra-Blade", "Inter-Blade", "Inter-Rack"]
        
        print(f"\n📊 消息大小: {message_size} 字节 (64KB)")
        print("通信类型\t\t预测延迟(μs)")
        print("-" * 40)
        
        predictions = []
        for i, comm_type in enumerate(comm_types):
            try:
                prediction = predict_ml_comm_time(message_size, comm_type=comm_type)
                predictions.append(prediction)
                print(f"{type_names[i]:<15}\t{prediction:>10.3f}")
                
                assert prediction > 0, f"预测值应该为正数，但得到: {prediction}"
                
            except Exception as e:
                print(f"❌ 预测通信类型 {comm_type} 时出错: {e}")
                return False
        
        # 计算平均值并与直接调用平均预测比较
        manual_avg = np.mean(predictions)
        direct_avg = predict_ml_comm_time(message_size, comm_type=None)
        
        print(f"\n手动计算平均值: {manual_avg:.3f} μs")
        print(f"直接预测平均值: {direct_avg:.3f} μs")
        print(f"差异: {abs(manual_avg - direct_avg):.6f} μs")
        
        # 验证两种方法的结果应该相近
        assert abs(manual_avg - direct_avg) < 0.001, "手动平均值与直接平均值差异过大"
        
        print("✅ 各通信类型预测测试通过")
        return True
        
    except Exception as e:
        print(f"❌ 通信类型预测测试失败: {e}")
        return False

def test_ml_analyze_function():
    """测试ML_analyze.py中的预测函数"""
    print("\n🔬 测试ML_analyze.py中的预测函数...")
    
    try:
        test_sizes = [1024, 16384, 65536]
        
        print("\n📊 测试get_ML_Prediction_results函数:")
        print("消息大小(字节)\t预测延迟(μs)")
        print("-" * 40)
        
        for size in test_sizes:
            try:
                prediction = get_ML_Prediction_results(size, comm_type=51)
                print(f"{size:>12}\t{prediction:>10.3f}")
                
                assert prediction > 0, f"预测值应该为正数，但得到: {prediction}"
                
            except Exception as e:
                print(f"❌ 测试get_ML_Prediction_results({size}, 51)时出错: {e}")
                return False
        
        print("✅ ML_analyze函数测试通过")
        return True
        
    except Exception as e:
        print(f"❌ ML_analyze函数测试失败: {e}")
        return False

def test_model_loading():
    """测试模型加载"""
    print("\n🔬 测试模型加载...")
    
    try:
        predictor = get_ml_predictor()
        
        if predictor.model is not None:
            print("✅ 成功加载ML模型")
            print(f"模型路径: {predictor.model_path}")
        else:
            print("⚠️  ML模型未加载，将使用备用预测方法")
            print(f"尝试的模型路径: {predictor.model_path}")
        
        return True
        
    except Exception as e:
        print(f"❌ 模型加载测试失败: {e}")
        return False

def main():
    """主测试函数"""
    print("🚀 开始ML预测功能测试")
    print("=" * 50)
    
    tests = [
        ("模型加载测试", test_model_loading),
        ("ML预测器基本功能", test_ml_predictor),
        ("各通信类型预测", test_individual_comm_types),
        ("ML_analyze函数", test_ml_analyze_function),
    ]
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        print(f"\n{'='*20} {test_name} {'='*20}")
        try:
            if test_func():
                passed += 1
                print(f"✅ {test_name} 通过")
            else:
                print(f"❌ {test_name} 失败")
        except Exception as e:
            print(f"❌ {test_name} 异常: {e}")
    
    print(f"\n{'='*50}")
    print(f"🎯 测试结果: {passed}/{total} 通过")
    
    if passed == total:
        print("🎉 所有测试通过！ML预测功能正常工作")
        return True
    else:
        print("⚠️  部分测试失败，请检查配置和模型文件")
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)